package services

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"

	"github.com/ollama/ollama/api"
	"net/url"
)

type OllamaService struct {
	client       *api.Client
	currentModel string
	knowledgeDir string
	mu           sync.RWMutex
}

func NewOllamaService(baseURL string) (*OllamaService, error) {
	client := api.NewClient(&url.URL{
		Scheme: "http",
		Host:   baseURL,
	}, nil)
	return &OllamaService{
		client:       client,
		knowledgeDir: "./knowledge",
	}, nil
}

func (s *OllamaService) SwitchModel(modelName string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.currentModel = modelName
	return nil
}

func (s *OllamaService) Query(ctx context.Context, prompt string, data interface{}) (string, error) {
	req := &api.GenerateRequest{
		Model:  s.currentModel,
		Prompt: prompt,
	}
	var result string
	callback := func(response api.GenerateResponse) error {
		result = response.Response
		return nil
	}
	err := s.client.Generate(ctx, req, callback)
	return result, err
}

func (s *OllamaService) storeKnowledge(data interface{}, result string) error {
	knowledge := map[string]interface{}{
		"data":   data,
		"result": result,
	}

	filePath := filepath.Join(s.knowledgeDir, fmt.Sprintf("%s.json", s.currentModel))
	file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer file.Close()

	encoder := json.NewEncoder(file)
	return encoder.Encode(knowledge)
}

func OllamaInit(baseURL string) (*OllamaService, error) {
	return NewOllamaService(baseURL)
}
